import csv
import numpy as np
import codecs

from utils import *

from sklearn.neighbors import KNeighborsClassifier

trainData,trainLabels,validData,validLabels=getTrainValid()

m1valid=codecs.open('./m1_valid.txt','w','UTF-8')
maxAcc=0
bestK=-1

accList=[]

for i in range(16,20):
    knn=KNeighborsClassifier(n_neighbors=i)
    knn.fit(trainData,trainLabels)
    predicted=knn.predict(validData)
    accuracy=(np.array(predicted)==validLabels).sum()*1.0/len(validLabels)
    accStr=str(i)+":"+str(accuracy)
    print(accStr)
    m1valid.write(accStr+'\n')
    accList.append(accuracy)
    if accuracy>maxAcc:
        bestK=i
        maxAcc=accuracy

plt.plot([i for i in range(16,20)],accList)
plt.show()

m1valid.write(str(bestK)+':'+str(maxAcc))
m1valid.close()
print(bestK,maxAcc)

knn=KNeighborsClassifier(n_neighbors=bestK)
knn.fit(trainData,trainLabels)

testFiles,testData=getTest()
testPredicted=knn.predict(testData)

csvHeaders=['id','categories']
csvData=[[testFiles[i],testPredicted[i]] for i in range(len(testFiles))]

with open('method1.csv', 'w', encoding='utf-8', newline='') as f:
    writer = csv.writer(f)
    writer.writerow(csvHeaders)
    writer.writerows(csvData)
